import torch
import torch.nn as nn

from enum import Enum
from typing import Any, Dict


class ModelType(Enum):
    DEEPLAB_V3 = 'deeplab_v3'
    DEEPLAB_V3_PLUS = 'deeplab_v3_plus'

    @classmethod
    def from_str(cls, label: str) -> "ModelType":
        if label in cls.__members__:
            return cls[label]
        
        for member in cls:
            if member.value.lower() == label.lower():
                return member
        raise ValueError(f"Unknown ModelType: {label!r}. "
                         f"Valid names: {list(cls.__members__.keys())}, "
                         f"values: {[m.value for m in cls]}")


def get_model(cfg: Dict[str, Any]):
    model_type = cfg.get('model')
    if isinstance(model_type, str):
        model_type = ModelType.from_str(model_type)
    model = None

    if model_type == ModelType.DEEPLAB_V3:
        from torchvision import models
        from torchvision.models.segmentation import DeepLabV3_ResNet101_Weights
        model = models.segmentation.deeplabv3_resnet101(weights=DeepLabV3_ResNet101_Weights.DEFAULT)
        model.classifier[-1] = nn.Conv2d(256, cfg.get('num_classes'), kernel_size=1)
    elif model_type == ModelType.DEEPLAB_V3_PLUS:
        from .deeplabV3Plus import DeeplabV3Plus
        model = DeeplabV3Plus(cfg)
    
    if torch.cuda.device_count() > 1:
        model = nn.DataParallel(model)
        
    return model
